# coding: UTF-8
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BertModel


class TextCNN(nn.Module):
    def __init__(self, config):
        super(TextCNN, self).__init__()
        self.bert = BertModel.from_pretrained(config.bert_path)
        for param in self.bert.parameters():
            param.requires_grad = True
        self.num_filters = config.num_filters
        self.filter_sizes = config.filter_sizes
        self.hidden_size = config.hidden_size
        self.conv1 = nn.Conv2d(1, self.num_filters, (self.filter_sizes[0], self.hidden_size))
        self.conv2 = nn.Conv2d(1, self.num_filters, (self.filter_sizes[1], self.hidden_size))
        self.conv3 = nn.Conv2d(1, self.num_filters, (self.filter_sizes[2], self.hidden_size))
        self.linear = nn.Linear(self.num_filters * 3, config.num_classes)

    def conv_and_pool(self, conv, x):
        out = conv(x)
        out = F.relu(out)
        return F.max_pool2d(out, (out.shape[2], out.shape[3])).squeeze()

    def forward(self, x, mask):
        out = self.bert(x, mask)[0].unsqueeze(1)
        out1 = self.conv_and_pool(self.conv1, out)
        out2 = self.conv_and_pool(self.conv2, out)
        out3 = self.conv_and_pool(self.conv3, out)
        out = torch.cat([out1, out2, out3], dim=1)
        return self.linear(out)


class SimpleModule(nn.Module):
    def __init__(self, config):
        super(SimpleModule, self).__init__()
        self.bert = BertModel.from_pretrained(config.bert_path)
        for param in self.bert.parameters():
            param.requires_grad = True
        self.num_filters = config.num_filters
        self.filter_sizes = config.filter_sizes
        self.hidden_size = config.hidden_size
        self.convs = nn.ModuleList([nn.Conv2d(1, self.num_filters, (i, self.hidden_size)) for i in self.filter_sizes])
        self.linear = nn.Linear(self.num_filters * 3, config.num_classes)

    def conv_and_pool(self, conv, x):
        out = conv(x)
        out = F.relu(out)
        return F.max_pool2d(out, (out.shape[2], out.shape[3])).squeeze()

    def forward(self, x, mask):
        out = self.bert(x, mask)[0].unsqueeze(1)
        out = torch.cat([self.conv_and_pool(conv, out) for conv in self.convs], dim=1)
        return self.linear(out)